package com.ofwiki.demo.scheduler.common;

import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.config.ScheduledTask;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
import org.springframework.scheduling.config.TriggerTask;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;

/**
 * @author HuangJS
 * @date 2021-12-11 3:04 下午
 */
public class DynamicCronScheduleTaskManager implements DisposableBean {
    private Map<Object, ScheduledTask> dynamicTaskMap = new HashMap<>();
    private ScheduledTaskRegistrar registrar;

    @Autowired
    protected DynamicCronScheduleTaskManager dynamicCronScheduleTaskManager;




    public DynamicCronScheduleTaskManager(ScheduledExecutorService scheduledExecutorService) {
        this.registrar = new ScheduledTaskRegistrar();
        this.registrar.setScheduler(scheduledExecutorService);
    }

    // 添加定时任务
    public ScheduledTask addTriggerTask(Object taskId, TriggerTask task) {
        ScheduledTask scheduledTask = dynamicTaskMap.get(taskId);
        if (scheduledTask != null) {
            scheduledTask.cancel();
        }

        scheduledTask = this.registrar.scheduleTriggerTask(task);
        dynamicTaskMap.put(taskId, scheduledTask);
        return scheduledTask;
    }
    
    public boolean contains(Object taskId){
        return this.dynamicTaskMap.containsKey(taskId);
    }
    
    // 更新定时任务的触发时机
    public void updateTriggerTask(Object taskId) {
        ScheduledTask scheduledTask = dynamicTaskMap.get(taskId);
        if (scheduledTask == null) {
            throw new IllegalStateException("Invalid taskId " + taskId + ",no fund ScheduledTask");
        }
        scheduledTask.cancel();

        scheduledTask = this.registrar.scheduleTriggerTask((TriggerTask) scheduledTask.getTask());
        dynamicTaskMap.put(taskId, scheduledTask);
    }

    // 移除定时任务
    public void removeTriggerTask(Object taskId) {
        ScheduledTask scheduledTask = dynamicTaskMap.remove(taskId);
        if (scheduledTask != null) {
            scheduledTask.cancel();
        }
    }

    @Override
    public void destroy() throws Exception {
        for (ScheduledTask value : dynamicTaskMap.values()) {
            value.cancel();
        }
        this.dynamicTaskMap.clear();
    }
}